import os.path
import re
from functools import lru_cache
import pandas as pd
from schema import db_connect
from project_path import get_project_path
from utils import logger
from utils.config_load import load_prompt
from utils.embedding import batch_embedding, single_embedding, cosine_similarity
from utils.llm import llm_api


def normalize_text(text:  str) -> str:
    """标准化文本，移除非字母数字字符并添加中文问号"""
    return re.sub(r'\W', '', text) + "？"


@lru_cache(maxsize=None)
def load_common_qa() -> tuple:
    """加载并预处理常见问答数据，使用缓存提升性能"""
    # 从数据库中，按照分类加载常见问题
    com_ques_data = db_connect.execute_sql("SELECT qa_level, qa_example, qa_sql FROM common_question")
    df = pd.DataFrame(list(com_ques_data.values())[0])
    df['processed_q'] = df['qa_example'].str.replace(r'\W', '', regex=True) + "？"
    # 常见问题层级字典
    layer_dict = {chr(65 + i): value for i, value in enumerate(df['qa_level'].drop_duplicates())}
    return df, dict(zip(df['processed_q'], df['qa_sql'])), layer_dict


@lru_cache(maxsize=None)
def load_statis_qa() -> pd.DataFrame:
    """加载并预处理统计问答数据，使用缓存提升性能"""
    statis_qa_path = os.path.join(get_project_path(), 'data', 'question_embedding.pkl')
    if os.path.exists(statis_qa_path):
        return pd.read_pickle(statis_qa_path)
    else:
        # 若统计问题已修改，则重新生成嵌入
        stat_ques_data = db_connect.execute_sql("SELECT stat_question, stat_num FROM stat_question")
        df = pd.DataFrame(list(stat_ques_data.values())[0])
        return batch_embedding(df)


def common_question(question: str) -> dict:
    """
    根据用户输入，匹配最相关的常见问题或常见问答对
    :param question: 用户输入
    :return: 常见问题 SQL
    """
    # 加载常见问题
    qa_df, qa_mapping, layer_dict = load_common_qa()
    # 标准化用户提问
    processed_q = normalize_text(question)
    # 精确匹配
    if sql := qa_mapping.get(processed_q, ''):
        logger.info(f"full match 【common question】 {question} 【sql】 {sql}")
        return {'success': sql}
    # LLM 语义匹配
    option = llm_api(load_prompt(question, layer_dict, 'commonQuestionMatch'))
    if layer := layer_dict.get(option, ''):
        layer_qa = qa_df[qa_df.qa_level == layer]
        qa_pair = dict(zip(layer_qa.qa_example, layer_qa.qa_sql))
        logger.info(f"partial match 【common question】 {layer} 【q-a】 {qa_pair}")
        return {'fail': qa_pair}
    else:
        logger.info("mismatch 【common question】")
        return {'fail': 'None'}


def statis_question(question: str) -> dict:
    """
    根据用户输入和统计表，匹配最相关问题的统计值
    :param question: 用户输入
    :return: 统计值
    """
    # LLM 判断是否为统计问题
    is_statis = llm_api(load_prompt(user_input=question, task='statisQuestionMatch'))
    if is_statis == "Yes":
        # 计算问题嵌入
        question_embedding = single_embedding(question)
        # 加载统计问题
        df = load_statis_qa()
        embeddings = df['question_embedding'].tolist()  # 问题嵌入
        df['cos_score'] = [cosine_similarity(question_embedding, embedding) for embedding in embeddings]
        # 根据 cos_score 降序排列 df
        df = df.sort_values(by='cos_score', ascending=False)
        if df.iloc[0]['cos_score'] > 0.96:
            stat_num = df.iloc[0]['stat_num']
            logger.info(f"match 【statis question】 {stat_num}")
            return {'success': stat_num}
        else:
            logger.info("mismatch 【statis question】")
            return {'fail': None}
    else:
        logger.info("mismatch 【statis question】")
        return {'fail': None}


if __name__ == '__main__':
    print(common_question("your question"))
    statis_question("your question")
